/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.io;



import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.io.DefaultFilenamePolicy.Params;

import com.bff.gaia.unified.sdk.io.DefaultFilenamePolicy.ParamsCoder;

import com.bff.gaia.unified.sdk.io.FileBasedSink.DynamicDestinations;

import com.bff.gaia.unified.sdk.io.FileBasedSink.FilenamePolicy;

import com.bff.gaia.unified.sdk.transforms.SerializableFunction;

import com.bff.gaia.unified.sdk.transforms.SerializableFunctions;

import com.bff.gaia.unified.sdk.transforms.display.DisplayData;



import javax.annotation.Nullable;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkState;



/** Some helper classes that derive from {@link FileBasedSink.DynamicDestinations}. */

public class DynamicFileDestinations {

  /** Always returns a constant {@link FilenamePolicy}. */

  private static class ConstantFilenamePolicy<UserT, OutputT>

      extends DynamicDestinations<UserT, Void, OutputT> {

    private final FilenamePolicy filenamePolicy;

    private final SerializableFunction<UserT, OutputT> formatFunction;



    public ConstantFilenamePolicy(

		FilenamePolicy filenamePolicy, SerializableFunction<UserT, OutputT> formatFunction) {

      this.filenamePolicy = filenamePolicy;

      this.formatFunction = formatFunction;

    }



    @Override

    public OutputT formatRecord(UserT record) {

      return formatFunction.apply(record);

    }



    @Override

    @Nullable

    public Void getDestination(UserT element) {

      return (Void) null;

    }



    @Override

    @Nullable

    public Coder<Void> getDestinationCoder() {

      return null;

    }



    @Override

    @Nullable

    public Void getDefaultDestination() {

      return (Void) null;

    }



    @Override

    public FilenamePolicy getFilenamePolicy(Void destination) {

      return filenamePolicy;

    }



    @Override

    public void populateDisplayData(DisplayData.Builder builder) {

      checkState(filenamePolicy != null);

      filenamePolicy.populateDisplayData(builder);

    }

  }



  /**

   * A base class for a {@link DynamicDestinations} object that returns differently-configured

   * instances of {@link DefaultFilenamePolicy}.

   */

  private static class DefaultPolicyDestinations<UserT, OutputT>

      extends DynamicDestinations<UserT, Params, OutputT> {

    private final SerializableFunction<UserT, Params> destinationFunction;

    private final Params emptyDestination;

    private final SerializableFunction<UserT, OutputT> formatFunction;



    public DefaultPolicyDestinations(

        SerializableFunction<UserT, Params> destinationFunction,

        Params emptyDestination,

        SerializableFunction<UserT, OutputT> formatFunction) {

      this.destinationFunction = destinationFunction;

      this.emptyDestination = emptyDestination;

      this.formatFunction = formatFunction;

    }



    @Override

    public OutputT formatRecord(UserT record) {

      return formatFunction.apply(record);

    }



    @Override

    public Params getDestination(UserT element) {

      return destinationFunction.apply(element);

    }



    @Override

    public Params getDefaultDestination() {

      return emptyDestination;

    }



    @Nullable

    @Override

    public Coder<Params> getDestinationCoder() {

      return ParamsCoder.of();

    }



    @Override

    public FilenamePolicy getFilenamePolicy(DefaultFilenamePolicy.Params params) {

      return DefaultFilenamePolicy.fromParams(params);

    }

  }



  /** Returns a {@link DynamicDestinations} that always returns the same {@link FilenamePolicy}. */

  public static <UserT, OutputT> DynamicDestinations<UserT, Void, OutputT> constant(

	  FilenamePolicy filenamePolicy, SerializableFunction<UserT, OutputT> formatFunction) {

    return new ConstantFilenamePolicy<>(filenamePolicy, formatFunction);

  }



  /**

   * A specialization of {@link #constant(FilenamePolicy, SerializableFunction)} for the case where

   * UserT and OutputT are the same type and the format function is the identity.

   */

  public static <UserT> DynamicDestinations<UserT, Void, UserT> constant(

      FilenamePolicy filenamePolicy) {

    return new ConstantFilenamePolicy<>(filenamePolicy, SerializableFunctions.<UserT>identity());

  }



  /**

   * Returns a {@link DynamicDestinations} that returns instances of {@link DefaultFilenamePolicy}

   * configured with the given {@link Params}.

   */

  public static <UserT, OutputT> DynamicDestinations<UserT, Params, OutputT> toDefaultPolicies(

      SerializableFunction<UserT, Params> destinationFunction,

      Params emptyDestination,

      SerializableFunction<UserT, OutputT> formatFunction) {

    return new DefaultPolicyDestinations<>(destinationFunction, emptyDestination, formatFunction);

  }

}